import sys, os
import keras
import cv2, shutil
import numpy as np

from src.keras_utils import load_model
from glob import glob
from os.path import splitext, basename
from src.utils import im2single
from src.keras_utils import load_model, detect_lp
from src.label import Shape, writeShapes
import time
from src.drawing_utils import draw_losangle


def adjust_pts(pts, lroi):
    return pts * lroi.wh().reshape((2, 1)) + lroi.tl().reshape((2, 1))


if __name__ == '__main__':

    input_dir = sys.argv[1]
    output_dir = input_dir

    lp_threshold = .5

    wpod_net_path = sys.argv[2]
    wpod_net = load_model(wpod_net_path)

    imgs_paths = glob('%s/*.jpg' % input_dir)

    print('Searching for license plates using WPOD-NET')

    for i, img_path in enumerate(imgs_paths):
        start_full = time.time()
        print('\t Processing %s'.format(img_path))

        bname = splitext(basename(img_path))[0]
        Ivehicle = cv2.imread(img_path)
        im_h = Ivehicle.shape[0]
        im_w = Ivehicle.shape[1]
        ratio = float(max(Ivehicle.shape[:2])) / min(Ivehicle.shape[:2])
        side = int(ratio * 288.)
        bound_dim = min(side + (side % (2 ** 4)), 608)
        print("\t\tBound dim: %d, ratio: %f".format(bound_dim, ratio))

        start = time.time()
        Llp, LlpImgs, _ = detect_lp(wpod_net, im2single(Ivehicle), bound_dim, 2 ** 4, (240, 80), lp_threshold)
        print(" detecting cost : {}".format(time.time() - start))

        RED = (0, 0, 255)
        if len(LlpImgs):
            Ilp = LlpImgs[0]
            # Ilp = cv2.cvtColor(Ilp, cv2.COLOR_BGR2GRAY)
            # Ilp = cv2.cvtColor(Ilp, cv2.COLOR_GRAY2BGR)

            I = Ivehicle
            pts = Shape(Llp[0].pts).pts
            ptspx = pts * np.array(I.shape[1::-1], dtype=float).reshape(2, 1)
            draw_losangle(I, ptspx, RED, 3)
            cv2.imwrite('%s/%s_full.png' % (output_dir, bname), I)
            cv2.imwrite('%s/%s_lp.png' % (output_dir, bname), Ilp * 255.)
            # writeShapes('%s/%s_lp.txt' % (output_dir,bname),[s])
            det_path = os.path.join(input_dir, 'detected')
            if os.path.exists(det_path) is False:
                os.makedirs(det_path)
            shutil.move('%s/%s_full.png' % (output_dir, bname), '%s/%s_full.png' % (det_path, bname))
            shutil.move('%s/%s_lp.png' % (output_dir, bname), '%s/%s_lp.png' % (det_path, bname))
            shutil.move(img_path, '%s/%s' % (det_path, basename(img_path)))
        print(" detecting full cost : {}".format(time.time() - start_full))
